package leetcode

// https://leetcode.cn/problems/swim-in-rising-water/

import "sort"

func swimInWater(grid [][]int) int {
	n := len(grid)
	var points = make([][]int, n*n)
	for i := range points {
		points[i] = make([]int, 3)
	}
	pi := 0
	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			points[pi][0] = i
			points[pi][1] = j
			points[pi][2] = grid[i][j]
			pi++
		}
	}
	sort.Slice(points, func(i, j int) bool {
		return points[i][2] < points[j][2]
	})
	uf := NewPoolUnionFind(n)
	ans := 0
	for i := 0; i < len(points); i++ {
		r := points[i][0]
		c := points[i][1]
		v := points[i][2]
		if r > 0 && grid[r-1][c] <= v {
			uf.Union(r, c, r-1, c)
		}
		if r < n-1 && grid[r+1][c] <= v {
			uf.Union(r, c, r+1, c)
		}
		if c > 0 && grid[r][c-1] <= v {
			uf.Union(r, c, r, c-1)
		}
		if c < n-1 && grid[r][c+1] <= v {
			uf.Union(r, c, r, c+1)
		}
		if uf.IsSameSet(0, 0, n-1, n-1) {
			ans = v
			break
		}
	}
	return ans
}

type PoolUnionFind struct {
	parent            []int
	size              []int
	help              []int
	rowNum, columnNum int
}

func NewPoolUnionFind(n int) *PoolUnionFind {
	size := n * n
	uf := &PoolUnionFind{make([]int, size), make([]int, size), make([]int, size), n, n}
	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			index := uf.Index(i, j)
			uf.parent[index] = index
			uf.size[index] = 1
		}
	}
	return uf
}

func (uf *PoolUnionFind) Find(i int) int {
	helpIndex := 0
	for i != uf.parent[i] {
		uf.help[helpIndex] = i
		helpIndex++

		i = uf.parent[i]
	}

	for helpIndex > 0 {
		helpIndex--
		uf.parent[uf.help[helpIndex]] = i
	}

	return i
}

func (uf *PoolUnionFind) Union(r1, c1, r2, c2 int) {
	index1 := uf.Index(r1, c1)
	index2 := uf.Index(r2, c2)

	root1 := uf.Find(index1)
	root2 := uf.Find(index2)

	if root1 == root2 {
		return
	}

	if uf.size[root1] >= uf.size[root2] {
		uf.parent[root2] = root1
		uf.size[root1] += uf.size[root2]
	} else {
		uf.parent[root1] = root2
		uf.size[root2] += uf.size[root1]
	}
}

func (uf *PoolUnionFind) IsSameSet(r1, c1, r2, c2 int) bool {
	index1 := uf.Index(r1, c1)
	index2 := uf.Index(r2, c2)

	return uf.Find(index1) == uf.Find(index2)
}

func (uf *PoolUnionFind) Index(rowIndex, columnIndex int) int {
	return rowIndex*uf.rowNum + columnIndex
}
